﻿using System;
using System.IO;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Gibbed.IO;
using Gibbed.MassEffect3.PackageUnpack;
using AmaroK86.MassEffect3.ZlibBlock;
//using Gibbed.MassEffect3.PackageDecompress;
//using ICSharpCode.SharpZipLib.Zip.Compression.Streams;

namespace AmaroK86.MassEffect3.PCCHandler
{
    public class PCCFile
    {
        public PCCBase pccInfo;
        public readonly FileStream pccStream;

        public PCCFile(string pccFileName)
        {
            if (PCCBase.pccIsCompressed(pccFileName))
                PCCUnpack.Decompress(pccFileName);
            pccInfo = new PCCBase(pccFileName);

            pccStream = new FileStream(pccFileName, FileMode.Open, FileAccess.ReadWrite); 
        }

        public void ReplaceBytes(int pccOffset, byte[] bytesToReplace, int byteStart, int byteEnd)
        {
            if (byteStart < 0) byteStart = 0;
            if (byteEnd > bytesToReplace.Length) byteEnd = bytesToReplace.Length;
            if (pccOffset < 0) pccOffset = 0;
            int count = byteEnd - byteStart;
            if (pccOffset + count > pccStream.Length)
                throw new ArgumentOutOfRangeException();
            pccStream.Seek(pccOffset, SeekOrigin.Begin);
            pccStream.Write(bytesToReplace, byteStart, count);
        }

        public void Close()
        {
            if (pccStream != null)
            {
                pccStream.Close();
                pccStream.Dispose();
                GC.Collect();
                GC.WaitForPendingFinalizers();
            }
            PCCPack.Compress(pccInfo);
        }

        /*~PCCFile()
        {
            if (pccStream != null)
                pccStream.Close();
        }*/
    }

    #region PCCBase

    public class PCCBase
    {
        public string pccFileName;
        public uint namesOffset;
        public uint exportInfosOffset;
        public uint importInfosOffset;
        public List<string> names;
        public ExportInfo[] exportInfos;
        public ImportInfo[] importInfos;
        public uint exportMinOffset;
        public readonly uint maxBlockSize = 0x100000;
        public readonly uint maxSegmentSize = 0x20000;
        public bool isCompressed { get; private set; }

        //public bool isCompressed { get { return _isCompressed; } }

        public PCCBase(string pccFileName)
        {
            this.pccFileName = pccFileName;

            using (var input = File.OpenRead(pccFileName))
            {
                var magic = input.ReadValueU32(Endian.Little);
                if (magic != 0x9E2A83C1 &&
                    magic.Swap() != 0x9E2A83C1)
                {
                    throw new FormatException("not a pcc file");
                }
                var endian = magic == 0x9E2A83C1 ?
                    Endian.Little : Endian.Big;
                var encoding = endian == Endian.Little ?
                    Encoding.Unicode : Encoding.BigEndianUnicode;

                var versionLo = input.ReadValueU16(endian);
                var versionHi = input.ReadValueU16(endian);

                if (versionLo != 684 &&
                    versionHi != 194)
                {
                    throw new FormatException("unsupported pccversion");
                }

                input.Seek(4, SeekOrigin.Current);

                var folderNameLength = input.ReadValueS32(endian);
                var folderNameByteLength =
                    folderNameLength >= 0 ? folderNameLength : (-folderNameLength * 2);
                input.Seek(folderNameByteLength, SeekOrigin.Current);

                /*var packageFlagsOffset = input.Position;*/
                var packageFlags = input.ReadValueU32(endian);

                if ((packageFlags & 8) != 0)
                {
                    input.Seek(4, SeekOrigin.Current);
                }

                var nameCount = input.ReadValueU32(endian);
                namesOffset = input.ReadValueU32(endian);
                var exportCount = input.ReadValueU32(endian);
                exportInfosOffset = input.ReadValueU32(endian);
                var importCount = input.ReadValueU32(endian);
                importInfosOffset = input.ReadValueU32(endian);

                Stream data;
                if ((packageFlags & 0x02000000) == 0)
                {
                    isCompressed = false;
                    data = input;
                }
                else
                {
                    isCompressed = true;
                    input.Seek(36, SeekOrigin.Current);

                    var generationsCount = input.ReadValueU32(endian);
                    input.Seek(generationsCount * 12, SeekOrigin.Current);

                    input.Seek(20, SeekOrigin.Current);

                    var blockCount = input.ReadValueU32(endian);

                    var blockStream = new BlockStream(input);
                    for (int i = 0; i < blockCount; i++)
                    {
                        var uncompressedOffset = input.ReadValueU32(endian);
                        var uncompressedSize = input.ReadValueU32(endian);
                        var compressedOffset = input.ReadValueU32(endian);
                        var compressedSize = input.ReadValueU32(endian);
                        blockStream.AddBlock(
                            uncompressedOffset, uncompressedSize,
                            compressedOffset, compressedSize);
                    }

                    data = blockStream;
                }

                names = new List<string>((int)nameCount);

                exportInfos = new ExportInfo[exportCount];
                for (uint i = 0; i < exportCount; i++)
                {
                    exportInfos[i] = new ExportInfo();
                }

                importInfos = new ImportInfo[importCount];
                for (uint i = 0; i < importCount; i++)
                {
                    importInfos[i] = new ImportInfo();
                }

                data.Seek(namesOffset, SeekOrigin.Begin);
                for (uint i = 0; i < nameCount; i++)
                {
                    var nameLength = data.ReadValueS32(endian);

                    if (nameLength >= 0)
                    {
                        names.Add(data.ReadString(nameLength, true, Encoding.UTF8));
                    }
                    else
                    {
                        names.Add(data.ReadString(-nameLength * 2, true, encoding));
                    }
                }

                data.Seek(importInfosOffset, SeekOrigin.Begin);
                for (uint i = 0; i < importCount; i++)
                {
                    var importInfo = importInfos[i];

                    var packageNameIndex = data.ReadValueS32(endian);
                    importInfo.PackageName = names[packageNameIndex];

                    data.Seek(12, SeekOrigin.Current);

                    var outerIndex = data.ReadValueS32(endian);
                    importInfo.Outer = GetResource(exportInfos, importInfos, outerIndex);

                    var objectNameIndex = data.ReadValueS32(endian);
                    importInfo.ObjectName = names[objectNameIndex];

                    data.Seek(4, SeekOrigin.Current);
                }

                exportMinOffset = 0;
                data.Seek(exportInfosOffset, SeekOrigin.Begin);
                for (uint i = 0; i < exportCount; i++)
                {
                    var exportInfo = exportInfos[i];

                    exportInfo.PackageName = Path.GetFileNameWithoutExtension(pccFileName);

                    var classIndex = data.ReadValueS32(endian);
                    exportInfo.Class = GetResource(exportInfos, importInfos, classIndex);

                    data.Seek(4, SeekOrigin.Current);

                    var outerIndex = data.ReadValueS32(endian);
                    exportInfo.Outer = GetResource(exportInfos, importInfos, outerIndex);

                    var objectNameIndex = data.ReadValueS32(endian);
                    exportInfo.ObjectName = names[objectNameIndex];

                    data.Seek(16, SeekOrigin.Current);

                    exportInfo.DataSize = data.ReadValueU32(endian);
                    exportInfo.DataOffset = data.ReadValueU32(endian);
                    if (exportMinOffset == 0)
                    {
                        exportMinOffset = exportInfo.DataOffset;
                    }
                    if (exportInfo.DataOffset < exportMinOffset)
                    {
                        exportMinOffset = exportInfo.DataOffset;
                    }

                    data.Seek(4, SeekOrigin.Current);
                    var count = data.ReadValueU32(endian);
                    data.Seek(count * 4, SeekOrigin.Current);
                    data.Seek(20, SeekOrigin.Current);
                }

                /*uint countSize = exportMinOffset - namesOffset;
                uint totalSize = exportMinOffset;
                for (int i = 0; i < exportInfos.Length; i++)
                {
                    ExportInfo exportInfo = exportInfos[i];
                    if (countSize + exportInfo.DataSize > maxBlockSize)
                    {
                        Console.WriteLine("Count size: {0} bytes, 0x{0:X8}, last data size: {1}", countSize, exportInfo.DataSize);
                        countSize = exportInfo.DataSize;
                    }
                    else
                    {
                        countSize += exportInfo.DataSize;
                    }
                    totalSize += exportInfo.DataSize;
                }*/

            }
        }

        private static IResource GetResource(ExportInfo[] exportInfos, ImportInfo[] importInfos, int index)
        {
            if (index == 0)
            {
                return null;
            }
            if (index < 0 && -index <= importInfos.Length)
            {
                return importInfos[-index - 1];
            }

            if (index > 0 && index <= exportInfos.Length)
            {
                return exportInfos[index - 1];
            }

            throw new InvalidOperationException();
        }

        public static bool pccIsCompressed(string pccFileName)
        {
            using (var input = File.OpenRead(pccFileName))
            {
                var magic = input.ReadValueU32(Endian.Little);
                if (magic != 0x9E2A83C1 &&
                    magic.Swap() != 0x9E2A83C1)
                {
                    throw new FormatException("not a pcc file");
                }
                var endian = magic == 0x9E2A83C1 ?
                    Endian.Little : Endian.Big;
                var encoding = endian == Endian.Little ?
                    Encoding.Unicode : Encoding.BigEndianUnicode;

                var versionLo = input.ReadValueU16(endian);
                var versionHi = input.ReadValueU16(endian);

                if (versionLo != 684 &&
                    versionHi != 194)
                {
                    throw new FormatException("unsupported pccversion");
                }

                input.Seek(4, SeekOrigin.Current);

                var folderNameLength = input.ReadValueS32(endian);
                var folderNameByteLength =
                    folderNameLength >= 0 ? folderNameLength : (-folderNameLength * 2);
                input.Seek(folderNameByteLength, SeekOrigin.Current);

                var packageFlags = input.ReadValueU32(endian);
                input.Close();
                return ((packageFlags & 0x02000000) != 0);
            }
        }

        
        public Stream ToStream()
        {
            MemoryStream mStream;
            if (isCompressed)
                mStream = (MemoryStream)PCCUnpack.DecompressToStream(this);
            else
            {
                mStream = new MemoryStream();
                FileStream oStream = new FileStream(pccFileName, FileMode.Open, FileAccess.Read);
                oStream.CopyTo(mStream);
                oStream.Close();
            }
            return mStream;
        }

        public byte[] ToArray()
        {
            FileStream oStream = new FileStream(pccFileName, FileMode.Open, FileAccess.Read);
            MemoryStream mStream = new MemoryStream();
            oStream.CopyTo(mStream);
            oStream.Close();
            return mStream.ToArray();
        }
    } //end PCCBase

    #endregion

    #region PCCPack

    public static class PCCPack
    {
        public static void Compress(string pccFileName, bool makeBackup = false)
        {
            PCCBase pccBase = new PCCBase(pccFileName);
            Compress(pccBase, makeBackup);
        }

        public static void Compress(PCCBase pccBase, bool makeBackup = false)
        {
            if (pccBase.isCompressed)
            {
                throw new FormatException("pcc file is already compressed");
            }

            string pccTempFile = pccBase.pccFileName + ".bak";
            if (File.Exists(pccTempFile))
                File.Delete(pccTempFile);
            File.Move(pccBase.pccFileName, pccTempFile);

            using (FileStream inputStream = File.OpenRead(pccTempFile), outputStream = File.OpenWrite(pccBase.pccFileName))
            {
                var magic = inputStream.ReadValueU32(Endian.Little);
                if (magic != 0x9E2A83C1 &&
                    magic.Swap() != 0x9E2A83C1)
                {
                    throw new FormatException("not a pcc file");
                }
                var endian = magic == 0x9E2A83C1 ? Endian.Little : Endian.Big;

                var versionLo = inputStream.ReadValueU16(endian);
                var versionHi = inputStream.ReadValueU16(endian);

                if (versionLo != 684 &&
                    versionHi != 194)
                {
                    throw new FormatException("unsupported pcc version");
                }

                byte[] buffer = new byte[130];
                inputStream.Seek(0, SeekOrigin.Begin);
                inputStream.Read(buffer, 0, 130);
                outputStream.Write(buffer, 0, buffer.Length);

                //add compressed pcc flag
                inputStream.Seek(12, SeekOrigin.Begin);
                var folderNameLength = inputStream.ReadValueS32();
                var folderNameByteLength =
                    folderNameLength >= 0 ? folderNameLength : (-folderNameLength * 2);
                inputStream.Seek(folderNameByteLength, SeekOrigin.Current);
                outputStream.Seek(inputStream.Position, SeekOrigin.Begin);

                var packageFlags = inputStream.ReadValueU32();
                if ((packageFlags & 0x02000000u) != 0)
                {
                    throw new FormatException("pcc file is already compressed");
                }
                packageFlags |= 0x02000000; // add compression flag
                outputStream.WriteValueU32(packageFlags);

                outputStream.Seek(buffer.Length, SeekOrigin.Begin);

                long outOffsetData;
                long outOffsetBlockInfo;
                long inOffsetData = pccBase.namesOffset;
                List<uint> blockSizes = new List<uint>();
                uint countSize = pccBase.exportMinOffset - pccBase.namesOffset;

                //count the number of blocks and relative sizes
                for (int i = 0; i < pccBase.exportInfos.Length; i++)
                {
                    ExportInfo exportInfo = pccBase.exportInfos[i];
                    if (countSize + exportInfo.DataSize > pccBase.maxBlockSize)
                    {
                        blockSizes.Add(countSize);
                        countSize = exportInfo.DataSize;
                    }
                    else
                    {
                        countSize += exportInfo.DataSize;
                    }
                }
                blockSizes.Add(countSize);

                outputStream.WriteValueS32(blockSizes.Count);
                outOffsetBlockInfo = outputStream.Position;
                outOffsetData = pccBase.namesOffset + (blockSizes.Count * 16);

                inputStream.Seek(pccBase.namesOffset, SeekOrigin.Begin);
                //divide the block in segments
                for (int i = 0; i < blockSizes.Count; i++)
                {
                    outputStream.Seek(outOffsetBlockInfo, SeekOrigin.Begin);
                    outputStream.WriteValueU32((uint)inputStream.Position);
                    outputStream.WriteValueU32(blockSizes[i]);
                    outputStream.WriteValueU32((uint)outOffsetData);

                    byte[] inputBlock = new byte[blockSizes[i]];
                    inputStream.Read(inputBlock,0,(int)blockSizes[i]);
                    byte[] compressedBlock = ZBlock.Compress(inputBlock, 0, inputBlock.Length);

                    outputStream.WriteValueS32(compressedBlock.Length);
                    outOffsetBlockInfo = outputStream.Position;

                    outputStream.Seek(outOffsetData, SeekOrigin.Begin);
                    outputStream.Write(compressedBlock, 0, compressedBlock.Length);
                    outOffsetData = outputStream.Position;
                }

                //copying some unknown values
                buffer = new byte[8];
                inputStream.Seek(pccBase.namesOffset - 8, SeekOrigin.Begin);
                inputStream.Read(buffer, 0, buffer.Length);
                outputStream.Seek(outOffsetBlockInfo, SeekOrigin.Begin);
                outputStream.Write(buffer, 0, buffer.Length);
            }

            if (!makeBackup)
                File.Delete(pccTempFile);
        }
    }

    #endregion

    #region PCCUnpack

    public static class PCCUnpack
    {
        public static void Decompress(string pccFileName, bool makeBackup = false, string pccBackupName = "")
        {
            string pccBackupFile = pccBackupName == "" ? pccFileName + ".bak" : pccBackupName;
            string pccOriginalName = pccFileName;
            if (File.Exists(pccBackupFile))
                File.Delete(pccBackupFile);
            File.Move(pccFileName, pccBackupFile);

            var input = File.OpenRead(pccBackupFile);

            var magic = input.ReadValueU32(Endian.Little);
            if (magic != 0x9E2A83C1 &&
                magic.Swap() != 0x9E2A83C1)
            {
                input.Close();
                throw new FormatException("not a pcc file");
            }
            var endian = magic == 0x9E2A83C1 ? Endian.Little : Endian.Big;

            var versionLo = input.ReadValueU16(endian);
            var versionHi = input.ReadValueU16(endian);

            if (versionLo != 684 &&
                versionHi != 194)
            {
                input.Close();
                throw new FormatException("unsupported pcc version");
            }

            long headerSize = 8;

            input.Seek(4, SeekOrigin.Current);
            headerSize += 4;

            var folderNameLength = input.ReadValueS32(endian);
            headerSize += 4;

            var folderNameByteLength =
                folderNameLength >= 0 ? folderNameLength : (-folderNameLength * 2);
            input.Seek(folderNameByteLength, SeekOrigin.Current);
            headerSize += folderNameByteLength;

            var packageFlagsOffset = input.Position;
            var packageFlags = input.ReadValueU32(endian);
            headerSize += 4;

            if ((packageFlags & 0x02000000u) == 0)
            {
                input.Close();
                File.Move(pccBackupFile, pccFileName);
                //throw new FormatException("pcc file is already decompressed");
            }

            if ((packageFlags & 8) != 0)
            {
                input.Seek(4, SeekOrigin.Current);
                headerSize += 4;
            }

            input.Seek(60, SeekOrigin.Current);
            headerSize += 60;

            var generationsCount = input.ReadValueU32(endian);
            input.Seek(generationsCount * 12, SeekOrigin.Current);
            headerSize += generationsCount * 12;

            input.Seek(20, SeekOrigin.Current);
            headerSize += 24;

            var blockCount = input.ReadValueU32(endian);
            int headBlockOff = (int)input.Position;
            var afterBlockTableOffset = headBlockOff + (blockCount * 16);
            var indataOffset = afterBlockTableOffset + 8;
            byte[] buff;

            FileStream output = new FileStream(pccFileName, FileMode.Create, FileAccess.Write);

            input.Seek(0, SeekOrigin.Begin);
            output.Seek(0, SeekOrigin.Begin);

            output.WriteFromStream(input, headerSize);
            output.WriteValueU32(0, endian); // block count
            input.Seek(afterBlockTableOffset, SeekOrigin.Begin);
            output.WriteFromStream(input, 8);

            for (int i = 0; i < blockCount; i++)
            {
                input.Seek(headBlockOff, SeekOrigin.Begin);
                var uncompressedOffset = input.ReadValueU32(endian);
                var uncompressedSize = input.ReadValueU32(endian);
                var compressedOffset = input.ReadValueU32(endian);
                var compressedSize = input.ReadValueU32(endian);
                headBlockOff = (int)input.Position;

                buff = new byte[compressedSize];
                input.Seek(compressedOffset, SeekOrigin.Begin);
                input.Read(buff, 0, buff.Length);

                byte[] temp = ZBlock.Decompress(buff, 0, buff.Length);
                output.Seek(uncompressedOffset, SeekOrigin.Begin);
                output.Write(temp, 0, temp.Length);
            }

            output.Seek(packageFlagsOffset, SeekOrigin.Begin);
            output.WriteValueU32(packageFlags & ~0x02000000u, endian);
            output.Close();
            input.Close();

            if (!makeBackup)
                File.Delete(pccBackupFile);
        }

        public static Stream DecompressToStream(PCCBase pccBase)
        {
            if (!pccBase.isCompressed)
            {
                return pccBase.ToStream();
                //throw new FormatException("pcc file is already decompressed");
            }

            var input = File.OpenRead(pccBase.pccFileName);

            var magic = input.ReadValueU32(Endian.Little);
            if (magic != 0x9E2A83C1 &&
                magic.Swap() != 0x9E2A83C1)
            {
                throw new FormatException("not a pcc file");
            }
            var endian = magic == 0x9E2A83C1 ? Endian.Little : Endian.Big;

            var versionLo = input.ReadValueU16(endian);
            var versionHi = input.ReadValueU16(endian);

            if (versionLo != 684 &&
                versionHi != 194)
            {
                throw new FormatException("unsupported pcc version");
            }

            long headerSize = 8;

            input.Seek(4, SeekOrigin.Current);
            headerSize += 4;

            var folderNameLength = input.ReadValueS32(endian);
            headerSize += 4;

            var folderNameByteLength =
                folderNameLength >= 0 ? folderNameLength : (-folderNameLength * 2);
            input.Seek(folderNameByteLength, SeekOrigin.Current);
            headerSize += folderNameByteLength;

            var packageFlagsOffset = input.Position;
            var packageFlags = input.ReadValueU32(endian);
            headerSize += 4;

            if ((packageFlags & 0x02000000u) == 0)
            {
                throw new FormatException("pcc file is already decompressed");
            }

            if ((packageFlags & 8) != 0)
            {
                input.Seek(4, SeekOrigin.Current);
                headerSize += 4;
            }

            input.Seek(60, SeekOrigin.Current);
            headerSize += 60;

            var generationsCount = input.ReadValueU32(endian);
            input.Seek(generationsCount * 12, SeekOrigin.Current);
            headerSize += generationsCount * 12;

            input.Seek(20, SeekOrigin.Current);
            headerSize += 24;

            var blockCount = input.ReadValueU32(endian);
            int headBlockOff = (int)input.Position;
            var afterBlockTableOffset = headBlockOff + (blockCount * 16);
            var indataOffset = afterBlockTableOffset + 8;
            byte[] buff;

            Stream output = new MemoryStream();

            input.Seek(0, SeekOrigin.Begin);
            output.Seek(0, SeekOrigin.Begin);

            output.WriteFromStream(input, headerSize);
            output.WriteValueU32(0, endian); // block count
            input.Seek(afterBlockTableOffset, SeekOrigin.Begin);
            output.WriteFromStream(input, 8);

            for (int i = 0; i < blockCount; i++)
            {
                input.Seek(headBlockOff, SeekOrigin.Begin);
                var uncompressedOffset = input.ReadValueU32(endian);
                var uncompressedSize = input.ReadValueU32(endian);
                var compressedOffset = input.ReadValueU32(endian);
                var compressedSize = input.ReadValueU32(endian);
                headBlockOff = (int)input.Position;

                buff = new byte[compressedSize];
                input.Seek(compressedOffset, SeekOrigin.Begin);
                input.Read(buff, 0, buff.Length);

                byte[] temp = ZBlock.Decompress(buff, 0, buff.Length);
                output.Seek(uncompressedOffset, SeekOrigin.Begin);
                output.Write(temp, 0, temp.Length);
            }

            output.Seek(packageFlagsOffset, SeekOrigin.Begin);
            output.WriteValueU32(packageFlags & ~0x02000000u, endian);
            return output;
        } // end Decompress()
    }

    #endregion
}
